# Copyright (c) 2023, AgiBot Inc.
# All rights reserved.

import datetime
import threading
import time

import aimdk.GetNavigationState_aimrt_rpc_ros2
import aimdk.protocol_aimrt_rpc_pb2
import aimdk.protocol_pb2
import aimdk.SetNavigationTask_aimrt_rpc_ros2
import aimrt_py
# import planning_msgs.srv
import requests

# Global running flag
running = threading.Event()
running.set()


def get_navi_status():
    url = "http://127.0.0.1:39001/rpc/planning_msgs/srv/GetNavigationState"
    headers = {"Content-Type": "application/json"}
    data = {
        "header": {"stamp": {"sec": 0, "nanosec": 0}, "frame_id": ""},
    }
    rsp = requests.post(url, headers=headers, json=data)
    print(rsp)
    rsp.raise_for_status()
    return rsp.json()["state"]


class NavigationModule(aimrt_py.ModuleBase):
    def __init__(self):
        super().__init__()
        self.core = aimrt_py.CoreRef()
        self.logger = aimrt_py.LoggerRef()
        self.work_executor = aimrt_py.ExecutorRef()

        # RPC clients
        self.mc_action_proxy = None
        self.mapping_proxy = None
        self.localization_proxy = None
        self.get_navigation_state_proxy = None
        self.set_navigation_task_proxy = None

        self.navigation_task_running = False

    def Info(self):
        info = aimrt_py.ModuleInfo()
        info.name = "NavigationModule"
        return info

    def Initialize(self, core):
        assert isinstance(core, aimrt_py.CoreRef)
        self.core = core
        self.logger = self.core.GetLogger()

        aimrt_py.info(self.logger, "Navigation module initializing")

        try:
            # Get config
            module_cfg_file_path = self.core.GetConfigurator().GetConfigFilePath()
            aimrt_py.info(self.logger, f"Config file: {module_cfg_file_path}")

            # Get executor
            self.work_executor = self.core.GetExecutorManager().GetExecutor("work_thread_pool")
            if not self.work_executor:
                aimrt_py.error(self.logger, "Get executor failed")
                return False

            # Register RPC clients
            rpc_handle = self.core.GetRpcHandle()

            # McActionService
            self.mc_action_proxy = aimdk.protocol_aimrt_rpc_pb2.McActionServiceProxy(
                rpc_handle)
            if not self.mc_action_proxy.RegisterClientFunc(rpc_handle):
                aimrt_py.error(self.logger, "Register McActionService failed")
                return False

            # MappingService
            self.mapping_proxy = aimdk.protocol_aimrt_rpc_pb2.MappingServiceProxy(
                rpc_handle)
            if not self.mapping_proxy.RegisterClientFunc(rpc_handle):
                aimrt_py.error(self.logger, "Register MappingService failed")
                return False

            # LocalizationService
            self.localization_proxy = aimdk.protocol_aimrt_rpc_pb2.LocalizationServiceProxy(
                rpc_handle)
            if not self.localization_proxy.RegisterClientFunc(rpc_handle):
                aimrt_py.error(
                    self.logger, "Register LocalizationService failed")
                return False

            # GetNavigationState
            self.get_navigation_state_proxy = aimdk.GetNavigationState_aimrt_rpc_ros2.GetNavigationStateServiceProxy(
                rpc_handle
            )
            if not self.get_navigation_state_proxy.RegisterClientFunc(rpc_handle):
                aimrt_py.error(
                    self.logger, "Register GetNavigationState failed")
                return False

            # SetNavigationTask
            self.set_navigation_task_proxy = aimdk.SetNavigationTask_aimrt_rpc_ros2.SetNavigationTaskServiceProxy(
                rpc_handle
            )
            if not self.set_navigation_task_proxy.RegisterClientFunc(rpc_handle):
                aimrt_py.error(
                    self.logger, "Register SetNavigationTask failed")
                return False

        except Exception as e:
            aimrt_py.error(self.logger, f"Initialize failed: {e}")
            return False

        aimrt_py.info(self.logger, "Navigation module initialized")
        return True

    def Start(self):
        aimrt_py.info(self.logger, "Navigation module starting")

        try:
            # Start navigation task with delay
            def start_task():
                self.navigation_task_running = True
                try:
                    self.navigation_main_loop()
                    aimrt_py.info(
                        self.logger, "********************navigation_main_loop exit.")
                finally:
                    self.navigation_task_running = False
            self.work_executor.ExecuteAfter(
                datetime.timedelta(seconds=1), start_task)
            aimrt_py.info(self.logger, "After start task.")

        except Exception as e:
            aimrt_py.error(self.logger, f"Start failed: {e}")
            aimrt_py.info(self.logger, "Start failed, exit.")
            return False
        aimrt_py.info(self.logger, "Start finished.")
        return True

    def Shutdown(self):
        aimrt_py.info(self.logger, "Navigation module shutting down")
        global running
        running.clear()

        # Wait for task completion
        start_time = time.time()
        while self.navigation_task_running and time.time() - start_time < 3.0:
            time.sleep(0.1)

        aimrt_py.info(self.logger, "Navigation module shutdown completed")

    def navigation_main_loop(self):
        """Navigation main loop"""
        aimrt_py.info(self.logger, "Navigation main loop started")

        # Get map ID
        current_map_id = self.get_current_map_id()
        if current_map_id is None or not running.is_set():
            return

        # Get target point
        target_id = self.get_target_point_id(current_map_id)
        if target_id is None or not running.is_set():
            return

        # Get current action
        current_action = self.get_current_action()
        # if current_action is None or not running.is_set():
        #     return

        # Set navigation action if needed
        # if current_action != aimdk.protocol_pb2.McAction.McAction_RL_NAVIGATION_DEFAULT:
        #    if not self.set_navigation_action() or not running.is_set():
        #        return

        # Set navigation task
        if not self.set_navigation_task(current_map_id, target_id) or not running.is_set():
            return

        # Monitor navigation state
        self.monitor_navigation_state()

        aimrt_py.info(self.logger, "Navigation main loop completed")

    def get_current_map_id(self):
        """Get current map ID"""
        ctx = self.mapping_proxy.NewContextSharedPtr()
        ctx.SetTimeout(datetime.timedelta(seconds=1))

        req = aimdk.protocol_pb2.GetCurrentWorkingMapReq()
        req.command = aimdk.protocol_pb2.MappingCommand.MappingCommand_GET_CURRENT_WORKING_MAP

        status, rsp = self.mapping_proxy.GetCurrentWorkingMap(ctx, req)
        if not status.OK():
            aimrt_py.warn(
                self.logger, f"GetCurrentWorkingMap failed: {status.ToString()}")
            return None

        aimrt_py.info(self.logger, f"Current map ID: {rsp.data.map_id}")
        return rsp.data.map_id

    def get_target_point_id(self, current_map_id):
        """Get target point ID"""
        ctx = self.localization_proxy.NewContextSharedPtr()
        ctx.SetTimeout(datetime.timedelta(seconds=1))

        req = aimdk.protocol_pb2.GetTopoMsgsReq()
        req.command = aimdk.protocol_pb2.TopoCommand.TopoCommand_GET_TOPO_MSG
        req.map_id = current_map_id

        status, rsp = self.localization_proxy.GetTopoMsgs(ctx, req)

        if not status.OK():
            aimrt_py.warn(
                self.logger, f"GetTopoMsgs failed: {status.ToString()}")
            return None

        if not rsp.data.points:
            aimrt_py.warn(self.logger, "No topo points found")
            return None

        points_id_list = []
        for point in rsp.data.points:
            points_id_list.append(point.point_id)

        aimrt_py.info(self.logger, f"Target point IDs: {points_id_list}")
        aimrt_py.info(
            self.logger, f"Input the target id: {self.__target_id__}")
        user_input_target_id = self.__target_id__
        if user_input_target_id not in points_id_list:
            aimrt_py.warn(
                self.logger, f"Target id {user_input_target_id} is not in the points id list")
            return None

        return user_input_target_id

    def set_target_id(self, input_target_id):
        print(f"set target id: {input_target_id}")
        self.__target_id__ = input_target_id

    def get_current_action(self):
        """Get current motion control state"""
        ctx = self.mc_action_proxy.NewContextSharedPtr()
        ctx.SetTimeout(datetime.timedelta(seconds=1))

        req = aimdk.protocol_pb2.CommonRequest()
        status, rsp = self.mc_action_proxy.GetAction(ctx, req)

        if not status.OK():
            aimrt_py.warn(
                self.logger, f"GetAction failed: {status.ToString()}")
            return None

        # Only continue if in stand default or navigation state
        if (
            rsp.info.current_action != aimdk.protocol_pb2.McAction.McAction_RL_LOCOMOTION_DEFAULT
            and rsp.info.current_action != aimdk.protocol_pb2.McAction.McAction_RL_NAVIGATION_DEFAULT
        ):
            aimrt_py.warn(
                self.logger, "Current action is not in RL_LOCOMOTION_DEFAULT or RL_NAVIGATION_DEFAULT state")
            return None

        return rsp.info.current_action

    def set_navigation_action(self):
        """Switch to navigation state"""
        ctx = self.mc_action_proxy.NewContextSharedPtr()
        ctx.SetTimeout(datetime.timedelta(seconds=1))

        req = aimdk.protocol_pb2.McActionRequest()
        req.command.action = aimdk.protocol_pb2.McAction.McAction_RL_NAVIGATION_DEFAULT

        status, _ = self.mc_action_proxy.SetAction(ctx, req)

        if not status.OK():
            aimrt_py.warn(
                self.logger, f"SetAction failed: {status.ToString()}")
            return False

        aimrt_py.info(self.logger, "Set navigation action success")
        return True

    def set_navigation_task(self, current_map_id, target_id):
        """Set navigation task"""
        ctx = self.set_navigation_task_proxy.NewContextSharedPtr()
        ctx.SetTimeout(datetime.timedelta(seconds=1))

        req = planning_msgs.srv.SetNavigationTask.Request()
        req.map_id = current_map_id
        req.type = req.NAV_TYPE_NORMAL
        req.target_id = target_id

        status, rsp = self.set_navigation_task_proxy.SetNavigationTask(
            ctx, req)
        if not status.OK():
            aimrt_py.warn(
                self.logger, f"SetNavigationTask failed: {status.ToString()}")
            return False

        if not rsp.status:
            aimrt_py.warn(self.logger, f"SetNavigationTask failed: {rsp.msg}")
            return False

        aimrt_py.info(self.logger, "Set navigation task success")
        return True

    def monitor_navigation_state(self):
        """Monitor navigation state"""

        while running.is_set():
            # try:
            # Context cannot be reused so we need to create a new one each time
            # 等待导航任务完成
            while True:
                navi_status = get_navi_status()
                print(f"当前导航状态: {navi_status}")
                if navi_status == "STATE_ARRIVED":
                    print("导航任务完成")
                    break
                time.sleep(1)
            # except Exception as e:
            #     aimrt_py.error(self.logger, f"Monitor error: {e}")
            #     if running.is_set():
            #         time.sleep(0.5)
            #     else:
            #         break
